/*
* ***** BEGIN GPL LICENSE BLOCK *****
*
* This program is free software; you can redistribute it and/or
* modify it under the terms of the GNU General Public License
* as published by the Free Software Foundation; either version 2
* of the License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program; if not, write to the Free Software Foundation,
* Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
*
* The Original Code is Copyright (C) 2008 Blender Foundation.
* All rights reserved.
*
* Contributor(s): Blender Foundation
*
* ***** END GPL LICENSE BLOCK *****
*/

#include <vector>
#include <map>
#include <algorithm> // std::find

#include "ExportSettings.h"
#include "BCAnimationCurve.h"
#include "BCAnimationSampler.h"
#include "collada_utils.h"

extern "C" {
#include "BKE_action.h"
#include "BKE_constraint.h"
#include "BKE_key.h"
#include "BKE_main.h"
#include "BKE_library.h"
#include "BKE_material.h"
#include "BLI_listbase.h"
#include "DNA_anim_types.h"
#include "DNA_scene_types.h"
#include "DNA_key_types.h"
#include "DNA_constraint_types.h"
#include "ED_object.h"
}

static std::string EMPTY_STRING;
static BCAnimationCurveMap BCEmptyAnimationCurves;

BCAnimationSampler::BCAnimationSampler(BlenderContext &blender_context, BCObjectSet &object_set):
	blender_context(blender_context)
{
	BCObjectSet::iterator it;
	for (it = object_set.begin(); it != object_set.end(); ++it) {
		Object *ob = *it;
		add_object(ob);
	}
}

BCAnimationSampler::~BCAnimationSampler()
{
	BCAnimationObjectMap::iterator it;
	for (it = objects.begin(); it != objects.end(); ++it) {
		BCAnimation *animation = it->second;
		delete animation;
	}
}

void BCAnimationSampler::add_object(Object *ob)
{
	BCAnimation *animation = new BCAnimation(blender_context.get_context(), ob);
	objects[ob] = animation;

	initialize_keyframes(animation->frame_set, ob);
	initialize_curves(animation->curve_map, ob);
}

BCAnimationCurveMap *BCAnimationSampler::get_curves(Object *ob)
{
	BCAnimation &animation = *objects[ob];
	if (animation.curve_map.size() == 0)
		initialize_curves(animation.curve_map, ob);
	return &animation.curve_map;
}

static void get_sample_frames(BCFrameSet &sample_frames, int sampling_rate, bool keyframe_at_end, Scene *scene)
{
	sample_frames.clear();

	if (sampling_rate < 1)
		return; // no sample frames in this case

	float sfra = scene->r.sfra;
	float efra = scene->r.efra;

	int frame_index;
	for (frame_index = nearbyint(sfra); frame_index < efra; frame_index += sampling_rate) {
		sample_frames.insert(frame_index);
	}

	if (frame_index >= efra && keyframe_at_end)
	{
		sample_frames.insert(efra);
	}
}

static bool is_object_keyframe(Object *ob, int frame_index)
{
	return false;
}

static void add_keyframes_from(bAction *action, BCFrameSet &frameset)
{
	if (action) {
		FCurve *fcu = NULL;
		for (fcu = (FCurve *)action->curves.first; fcu; fcu = fcu->next) {
			BezTriple  *bezt = fcu->bezt;
			for (int i = 0; i < fcu->totvert; bezt++, i++) {
				int frame_index = nearbyint(bezt->vec[1][0]);
				frameset.insert(frame_index);
			}
		}
	}
}

void BCAnimationSampler::check_property_is_animated(BCAnimation &animation, float *ref, float *val, std::string data_path, int length)
{
	for (int array_index =0; array_index < length; ++array_index) {
		if (!bc_in_range(ref[length], val[length], 0.00001)) {
			BCCurveKey key(BC_ANIMATION_TYPE_OBJECT, data_path, array_index);
			BCAnimationCurveMap::iterator it = animation.curve_map.find(key);
			if (it == animation.curve_map.end()) {
				animation.curve_map[key] = new BCAnimationCurve(key, animation.get_reference());
			}
		}
	}
}

void BCAnimationSampler::update_animation_curves(BCAnimation &animation, BCSample &sample, Object *ob, int frame)
{
	BCAnimationCurveMap::iterator it;
	for (it = animation.curve_map.begin(); it != animation.curve_map.end(); ++it) {
		BCAnimationCurve *curve = it->second;
		if (curve->is_transform_curve()) {
			curve->add_value_from_matrix(sample, frame);
		}
		else {
			curve->add_value_from_rna(frame);
		}
	}
}

BCSample &BCAnimationSampler::sample_object(Object *ob, int frame_index, bool for_opensim)
{
	BCSample &ob_sample = sample_data.add(ob, frame_index);

	if (ob->type == OB_ARMATURE) {
		bPoseChannel *pchan;
		for (pchan = (bPoseChannel *)ob->pose->chanbase.first; pchan; pchan = pchan->next) {
			Bone *bone = pchan->bone;
			Matrix bmat;
			if (bc_bone_matrix_local_get(ob, bone, bmat, for_opensim)) {
				ob_sample.add_bone_matrix(bone, bmat);
			}
		}
	}
	return ob_sample;
}

void BCAnimationSampler::sample_scene(
	int sampling_rate,
	int keyframe_at_end,
	bool for_opensim,
	bool keep_keyframes,
	BC_export_animation_type export_animation_type)
{
	Scene *scene = blender_context.get_scene();
	BCFrameSet scene_sample_frames;
	get_sample_frames(scene_sample_frames, sampling_rate, keyframe_at_end, scene);
	BCFrameSet::iterator it;

	int startframe = scene->r.sfra;
	int endframe = scene->r.efra;

	for (int frame_index = startframe; frame_index <= endframe; ++frame_index) {
		/* Loop over all frames and decide for each frame if sampling is necessary */
		bool is_scene_sample_frame = false;
		bool needs_update = true;
		if (scene_sample_frames.find(frame_index) != scene_sample_frames.end()) {
			bc_update_scene(blender_context, frame_index);
			needs_update = false;
			is_scene_sample_frame = true;
		}

		bool needs_sampling = is_scene_sample_frame || keep_keyframes || export_animation_type == BC_ANIMATION_EXPORT_KEYS;
		if (!needs_sampling) {
			continue;
		}

		BCAnimationObjectMap::iterator obit;
		for (obit = objects.begin(); obit != objects.end(); ++obit) {
			Object *ob = obit->first;
			BCAnimation *animation = obit->second;
			BCFrameSet &object_keyframes = animation->frame_set;
			if (is_scene_sample_frame || object_keyframes.find(frame_index) != object_keyframes.end()) {

				if (needs_update) {
					bc_update_scene(blender_context, frame_index);
					needs_update = false;
				}

				BCSample &sample = sample_object(ob, frame_index, for_opensim);
				update_animation_curves(*animation, sample, ob, frame_index);
			}
		}
	}
}

bool BCAnimationSampler::is_animated_by_constraint(Object *ob, ListBase *conlist, std::set<Object *> &animated_objects)
{
	bConstraint *con;
	for (con = (bConstraint *)conlist->first; con; con = con->next) {
		ListBase targets = { NULL, NULL };

		const bConstraintTypeInfo *cti = BKE_constraint_typeinfo_get(con);

		if (!bc_validateConstraints(con))
			continue;

		if (cti && cti->get_constraint_targets) {
			bConstraintTarget *ct;
			Object *obtar;
			cti->get_constraint_targets(con, &targets);
			for (ct = (bConstraintTarget *)targets.first; ct; ct = ct->next) {
				obtar = ct->tar;
				if (obtar) {
					if (animated_objects.find(obtar) != animated_objects.end())
						return true;
				}
			}
		}
	}
	return false;
}

void BCAnimationSampler::find_depending_animated(std::set<Object *> &animated_objects, std::set<Object *> &candidates)
{
	bool found_more;
	do {
		found_more = false;
		std::set<Object *>::iterator it;
		for (it = candidates.begin(); it != candidates.end(); ++it) {
			Object *cob = *it;
			ListBase *conlist = get_active_constraints(cob);
			if (is_animated_by_constraint(cob, conlist, animated_objects)) {
				animated_objects.insert(cob);
				candidates.erase(cob);
				found_more = true;
				break;
			}
		}
	} while (found_more && candidates.size() > 0);
}

void BCAnimationSampler::get_animated_from_export_set(std::set<Object *> &animated_objects, LinkNode &export_set)
{
	/*
	Check if this object is animated. That is: Check if it has its own action, or

	- Check if it has constraints to other objects
	- at least one of the other objects is animated as well
	*/

	animated_objects.clear();
	std::set<Object *> static_objects;
	std::set<Object *> candidates;

	LinkNode *node;
	for (node = &export_set; node; node = node->next) {
		Object *cob = (Object *)node->link;
		if (bc_has_animations(cob)) {
			animated_objects.insert(cob);
		}
		else {
			ListBase conlist = cob->constraints;
			if (conlist.first)
				candidates.insert(cob);
		}
	}
	find_depending_animated(animated_objects, candidates);
}

void BCAnimationSampler::get_object_frames(BCFrames &frames, Object *ob)
{
	sample_data.get_frames(ob, frames);
}

void BCAnimationSampler::get_bone_frames(BCFrames &frames, Object *ob, Bone *bone)
{
	sample_data.get_frames(ob, bone, frames);
}

bool BCAnimationSampler::get_bone_samples(BCMatrixSampleMap &samples, Object *ob, Bone *bone)
{
	sample_data.get_matrices(ob, bone, samples);
	return bc_is_animated(samples);
}

bool BCAnimationSampler::get_object_samples(BCMatrixSampleMap &samples, Object *ob)
{
	sample_data.get_matrices(ob, samples);
	return bc_is_animated(samples);
}

#if 0
/*
   Add sampled values to FCurve
   If no FCurve exists, create a temporary FCurve;
   Note: The temporary FCurve will later be removed when the
   BCAnimationSampler is removed (by its destructor)

   curve: The curve to whioch the data is added
   matrices: The set of matrix values from where the data is taken
   animation_type BC_ANIMATION_EXPORT_SAMPLES: Use all matrix data
   animation_type BC_ANIMATION_EXPORT_KEYS: Only take data from matrices for keyframes
*/

void BCAnimationSampler::add_value_set(
	BCAnimationCurve &curve,
	BCFrameSampleMap &samples,
	BC_export_animation_type animation_type)
{
	int array_index = curve.get_array_index();
	const BC_animation_transform_type tm_type = curve.get_transform_type();

	BCFrameSampleMap::iterator it;
	for (it = samples.begin(); it != samples.end(); ++it) {
		const int frame_index = nearbyint(it->first);
		if (animation_type == BC_ANIMATION_EXPORT_SAMPLES || curve.is_keyframe(frame_index)) {

			const BCSample *sample = it->second;
			float val = 0;

			int subindex = curve.get_subindex();
			bool good;
			if (subindex == -1) {
				good = sample->get_value(tm_type, array_index, &val);
			}
			else {
				good = sample->get_value(tm_type, array_index, &val, subindex);
			}

			if (good) {
				curve.add_value(val, frame_index);
			}
		}
	}
	curve.remove_unused_keyframes();
	curve.calchandles();
}
#endif

void BCAnimationSampler::generate_transform(
	Object *ob,
    const BCCurveKey &key,
	BCAnimationCurveMap &curves)
{
	BCAnimationCurveMap::const_iterator it = curves.find(key);
	if (it == curves.end()) {
		curves[key] = new BCAnimationCurve(key, ob);
	}
}

void BCAnimationSampler::generate_transforms(
	Object *ob,
	const std::string prep,
	const BC_animation_type type,
	BCAnimationCurveMap &curves)
{
	generate_transform(ob, BCCurveKey(type, prep+"location", 0), curves);
	generate_transform(ob, BCCurveKey(type, prep+"location", 1), curves);
	generate_transform(ob, BCCurveKey(type, prep+"location", 2), curves);
	generate_transform(ob, BCCurveKey(type, prep+"rotation_euler", 0), curves);
	generate_transform(ob, BCCurveKey(type, prep+"rotation_euler", 1), curves);
	generate_transform(ob, BCCurveKey(type, prep+"rotation_euler", 2), curves);
	generate_transform(ob, BCCurveKey(type, prep+"scale", 0), curves);
	generate_transform(ob, BCCurveKey(type, prep+"scale", 1), curves);
	generate_transform(ob, BCCurveKey(type, prep+"scale", 2), curves);
}

void BCAnimationSampler::generate_transforms(Object *ob, Bone *bone, BCAnimationCurveMap &curves)
{
	std::string prep = "pose.bones[\"" + std::string(bone->name) + "\"].";
	generate_transforms(ob, prep, BC_ANIMATION_TYPE_BONE, curves);

	for (Bone *child = (Bone *)bone->childbase.first; child; child = child->next)
		generate_transforms(ob, child, curves);
}

/*
* Collect all keyframes from all animation curves related to the object
* The bc_get... functions check for NULL and correct object type
* The add_keyframes_from() function checks for NULL
*/
void BCAnimationSampler::initialize_keyframes(BCFrameSet &frameset, Object *ob)
{
	frameset.clear();
	add_keyframes_from(bc_getSceneObjectAction(ob), frameset);
	add_keyframes_from(bc_getSceneCameraAction(ob), frameset);
	add_keyframes_from(bc_getSceneLampAction(ob), frameset);

	for (int a = 0; a < ob->totcol; a++) {
		Material *ma = give_current_material(ob, a + 1);
		add_keyframes_from(bc_getSceneMaterialAction(ma), frameset);
	}
}

void BCAnimationSampler::initialize_curves(BCAnimationCurveMap &curves, Object *ob)
{
	BC_animation_type object_type = BC_ANIMATION_TYPE_OBJECT;

	bAction *action = bc_getSceneObjectAction(ob);
	if (action) {
		FCurve *fcu = (FCurve *)action->curves.first;

		for (; fcu; fcu = fcu->next) {
			object_type = BC_ANIMATION_TYPE_OBJECT;
			if (ob->type == OB_ARMATURE) {
				char *boneName = BLI_str_quoted_substrN(fcu->rna_path, "pose.bones[");
				if (boneName) {
					object_type = BC_ANIMATION_TYPE_BONE;
				}
			}

			/* Adding action curves on object */
			BCCurveKey key(object_type, fcu->rna_path, fcu->array_index);
			curves[key] = new BCAnimationCurve(key, ob, fcu);
		}
	}

	/* Add missing curves */
	object_type = BC_ANIMATION_TYPE_OBJECT;
	generate_transforms(ob, EMPTY_STRING, object_type, curves);
	if (ob->type == OB_ARMATURE) {
		bArmature *arm = (bArmature *)ob->data;
		for (Bone *root_bone = (Bone *)arm->bonebase.first; root_bone; root_bone = root_bone->next)
			generate_transforms(ob, root_bone, curves);
	}

	/* Add curves on Object->data actions */
	action = NULL;
	if (ob->type == OB_CAMERA) {
		action = bc_getSceneCameraAction(ob);
		object_type = BC_ANIMATION_TYPE_CAMERA;
	}
	else if (ob->type == OB_LAMP) {
		action = bc_getSceneLampAction(ob);
		object_type = BC_ANIMATION_TYPE_LIGHT;
	}

	if (action) {
		/* Add lamp action or Camera action */
		FCurve *fcu = (FCurve *)action->curves.first;
		for (; fcu; fcu = fcu->next) {
			BCCurveKey key(object_type, fcu->rna_path, fcu->array_index);
			curves[key] = new BCAnimationCurve(key, ob, fcu);
		}
	}

	/* Add curves on Object->material actions*/
	object_type = BC_ANIMATION_TYPE_MATERIAL;
	for (int a = 0; a < ob->totcol; a++) {
		/* Export Material parameter animations. */
		Material *ma = give_current_material(ob, a + 1);
		if (ma) {
			action = bc_getSceneMaterialAction(ma);
			if (action) {
				/* isMatAnim = true; */
				FCurve *fcu = (FCurve *)action->curves.first;
				for (; fcu; fcu = fcu->next) {
					BCCurveKey key(object_type, fcu->rna_path, fcu->array_index, a);
					curves[key] = new BCAnimationCurve(key, ob, fcu);
				}
			}
		}
	}
}

/* ==================================================================== */

BCSample &BCSampleFrame::add(Object *ob)
{
	BCSample *sample = new BCSample(ob);
	sampleMap[ob] = sample;
	return *sample;
}

/* Get the matrix for the given key, returns Unity when the key does not exist */
const BCSample *BCSampleFrame::get_sample(Object *ob) const
{
	BCSampleMap::const_iterator it = sampleMap.find(ob);
	if (it == sampleMap.end()) {
		return NULL;
	}
	return it->second;
}

const BCMatrix *BCSampleFrame::get_sample_matrix(Object *ob) const
{
	BCSampleMap::const_iterator it = sampleMap.find(ob);
	if (it == sampleMap.end()) {
		return NULL;
	}
	BCSample *sample = it->second;
	return &sample->get_matrix();
}

/* Get the matrix for the given Bone, returns Unity when the Objewct is not sampled */
const BCMatrix *BCSampleFrame::get_sample_matrix(Object *ob, Bone *bone) const
{
	BCSampleMap::const_iterator it = sampleMap.find(ob);
	if (it == sampleMap.end()) {
		return NULL;
	}

	BCSample *sample = it->second;
	const BCMatrix *bc_bone = sample->get_matrix(bone);
	return bc_bone;
}

/* Check if the key is in this BCSampleFrame */
const bool BCSampleFrame::has_sample_for(Object *ob) const
{
	return sampleMap.find(ob) != sampleMap.end();
}

/* Check if the Bone is in this BCSampleFrame */
const bool BCSampleFrame::has_sample_for(Object *ob, Bone *bone) const
{
	const BCMatrix *bc_bone = get_sample_matrix(ob, bone);
	return (bc_bone);
}

/* ==================================================================== */

BCSample &BCSampleFrameContainer::add(Object *ob, int frame_index)
{
	BCSampleFrame &frame = sample_frames[frame_index];
	return frame.add(ob);
}


/* ====================================================== */
/* Below are the getters which we need to export the data */
/* ====================================================== */

/* Return either the BCSampleFrame or NULL if frame does not exist*/
BCSampleFrame * BCSampleFrameContainer::get_frame(int frame_index)
{
	BCSampleFrameMap::iterator it = sample_frames.find(frame_index);
	BCSampleFrame *frame = (it == sample_frames.end()) ? NULL : &it->second;
	return frame;
}

/* Return a list of all frames that need to be sampled */
const int BCSampleFrameContainer::get_frames(std::vector<int> &frames) const
{
	frames.clear(); // safety;
	BCSampleFrameMap::const_iterator it;
	for (it = sample_frames.begin(); it != sample_frames.end(); ++it) {
		frames.push_back(it->first);
	}
	return frames.size();
}

const int BCSampleFrameContainer::get_frames(Object *ob, BCFrames &frames) const
{
	frames.clear(); // safety;
	BCSampleFrameMap::const_iterator it;
	for (it = sample_frames.begin(); it != sample_frames.end(); ++it) {
		const BCSampleFrame &frame = it->second;
		if (frame.has_sample_for(ob)) {
			frames.push_back(it->first);
		}
	}
	return frames.size();
}

const int BCSampleFrameContainer::get_frames(Object *ob, Bone *bone, BCFrames &frames) const
{
	frames.clear(); // safety;
	BCSampleFrameMap::const_iterator it;
	for (it = sample_frames.begin(); it != sample_frames.end(); ++it) {
		const BCSampleFrame &frame = it->second;
		if (frame.has_sample_for(ob, bone)) {
			frames.push_back(it->first);
		}
	}
	return frames.size();
}

const int BCSampleFrameContainer::get_samples(Object *ob, BCFrameSampleMap &samples) const
{
	samples.clear(); // safety;
	BCSampleFrameMap::const_iterator it;
	for (it = sample_frames.begin(); it != sample_frames.end(); ++it) {
		const BCSampleFrame &frame = it->second;
		const BCSample *sample = frame.get_sample(ob);
		if (sample) {
			samples[it->first] = sample;
		}
	}
	return samples.size();
}

const int BCSampleFrameContainer::get_matrices(Object *ob, BCMatrixSampleMap &samples) const
{
	samples.clear(); // safety;
	BCSampleFrameMap::const_iterator it;
	for (it = sample_frames.begin(); it != sample_frames.end(); ++it) {
		const BCSampleFrame &frame = it->second;
		const BCMatrix *matrix = frame.get_sample_matrix(ob);
		if (matrix) {
			samples[it->first] = matrix;
		}
	}
	return samples.size();
}

const int BCSampleFrameContainer::get_matrices(Object *ob, Bone *bone, BCMatrixSampleMap &samples) const
{
	samples.clear(); // safety;
	BCSampleFrameMap::const_iterator it;
	for (it = sample_frames.begin(); it != sample_frames.end(); ++it) {
		const BCSampleFrame &frame = it->second;
		const BCMatrix *sample = frame.get_sample_matrix(ob, bone);
		if (sample) {
			samples[it->first] = sample;
		}
	}
	return samples.size();
}
